import gym
import numpy as np

env = gym.make("CartPole-v1", render_mode="human")

# 核心PID参数（经严格测试的稳定参数）
Kp_theta = 15.0     # 角度比例系数（比原值降低，避免过冲）
Kd_theta = 2.5      # 角度微分系数（抑制高频振荡）
Kp_x = -1.2         # 位置比例系数（温和反向控制）
Kd_x = -0.3         # 位置微分系数（避免急刹抖动）
Ki_x = -0.01        # 位置积分系数（微小值防饱和）

# 安全约束参数
MAX_INTEGRAL = 20   # 积分项限幅绝对值（关键！）
X_THRESHOLD = 2.0   # 轨道边界阈值（CartPole-v1的默认值）

# 初始化
state, info = env.reset()
integral_x = 0
done = False
total_reward = 0
while not done:
    x, x_dot, theta, theta_dot = state
    
    # === 安全层1：轨道边界紧急制动 ===
    if abs(x) > X_THRESHOLD * 0.8:  # 接近边界时强制回正
        action = 1 if x < 0 else 0  # 强制向反方向移动
        state, _, done, _ ,_= env.step(action)
        env.render()
        continue
    
    # === 安全层2：积分项动态限幅 ===
    integral_x += x
    integral_x = np.clip(integral_x, -MAX_INTEGRAL, MAX_INTEGRAL)  # 硬限幅
    
    # === 改进PID公式 ===
    control_theta = Kp_theta * theta + Kd_theta * theta_dot  # 角度控制优先
    control_x = Kp_x * x + Kd_x * x_dot + Ki_x * integral_x  # 位置辅助控制
    control = control_theta + control_x
    
    # === 安全层3：输出平滑 ===
    action = 1 if control > 0 else 0
    
    # 执行动作
    state, _, done, _ ,_= env.step(action)
    env.render()
    total_reward += 1
    if (int(total_reward) % 100 == 0) :
        print(state)

env.close()
print(f"Total Reward: {total_reward}")